#######################################
# 训练脚本 run.py
# 注意不是功能源码，仅是AI训练源码，供参考
# 可接受的参数：
# 数据生成 coco
# 模型训练 train  (单卡)
# 模型训练 train2 (双卡)
# 模型评估 eval
# 模型导出 export
# 模型转换 onnx
# 模型测试 infer (pd模型测试)
# 模型测试 infer2 (onnx模型测试)
#######################################

import sys
import os
import paddle2onnx

if __name__ == "__main__":
  if len(sys.argv) > 1:
    cmd = sys.argv[1]
    os.system("cls")
    print("start run " + cmd + " ...")
  else:
    cmd = "err"

  # 数据生成
  if cmd == "coco":
    os.system("python tools/x2coco.py  \
               --dataset_type labelme  \
               --json_input_dir D:/dna_data2/json/  \
               --image_input_dir D:/dna_data2/pic/  \
               --output_dir D:/dna_coco2  \
               --train_proportion 1.0  \
               --val_proportion 0.0  \
               --test_proportion 0.0")
  # 模型训练（单卡）
  elif cmd == "train":
    os.system("python tools/train.py  \
               -c configs/mask_rcnn/mask_rcnn_r50_vd_fpn_2x_coco.yml")  
  # 模型训练（双卡）  
  elif cmd == "train2":
    os.system("python -m paddle.distributed.launch  \
               --gpus 0,1 tools/train.py  \
               -c configs/mask_rcnn/mask_rcnn_r50_vd_fpn_1x_coco.yml")
  # 模型评估
  elif cmd == "eval":
    os.system("python tools/eval.py  \
               -c configs/mask_rcnn/mask_rcnn_r50_vd_fpn_2x_coco.yml  \
               -o weights=output/mask_rcnn_r50_vd_fpn_2x_coco/model_final.pdparams")
  # 模型导出
  elif cmd =="export":
    os.system("python tools/export_model.py  \
               -c configs/mask_rcnn/mask_rcnn_r50_vd_fpn_2x_coco.yml  \
               --output_dir=D:/DeveFolder/DNAobject/ckas_model  \
               -o weights=output/mask_rcnn_r50_vd_fpn_2x_coco/model_final.pdparams \
               export_onnx=True")
  # 生成onnx
  elif cmd =="onnx":
    os.system("paddle2onnx  \
               --model_dir D:/DeveFolder/DNAobject/ckas_model/ch_first  \
               --model_filename model.pdmodel  \
               --params_filename model.pdiparams  \
               --opset_version 16  \
               --enable_dev_version False \
               --save_file D:/DeveFolder/DNAobject/ckas_model/ch_first.onnx")  
  # 模型部署/测试
  elif cmd =="infer":
    os.system("python deploy/python/infer.py  \
               --model_dir=D:/DeveFolder/DNAobject/ckas_model/ch_first  \
               --image_dir=D:/DeveFolder/DNAobject/ckas_model/demo  \
               --device=GPU  \
               --output_dir=D:/DeveFolder/DNAobject/caks_model/output")
  # 模型部署/测试: onnx
  elif cmd == "infer2":
    os.system("python deploy/third_engine/onnx/infer.py  \
              --infer_cfg D:/DeveFolder/DNAobject/ckas_model/ch_first/infer_cfg.yml \
              --onnx_file D:/DeveFolder/DNAobject/ckas_model/ch_first.onnx  \
              --image_dir D:/DeveFolder/DNAobject/ckas_model/demo/")
  # 提示信息
  else:
    print("[信息]:参数只能是coco、train、train2、evel、export、infer, 其它参数无效")

